package com.hospital.diagnosis.module.hanlder;

import com.hospital.common.base.ThreadAttributes;
import com.hospital.common.utils.ThreadLocalUtils;
import org.springframework.http.HttpHeaders;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.util.ObjectUtils;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import java.util.List;
import java.util.Map;

import static com.hospital.common.base.interfaces.Const.*;

@Component
public class WebSocketHandshakeInterceptor implements HandshakeInterceptor {

    /**
     * websocket连接建立前处理,用于获取子协议token,请求头
     */
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
                                   WebSocketHandler wsHandler, Map<String, Object> attributes) {
        HttpHeaders headers = request.getHeaders();
        String userId = headers.getFirst(USER_ID_HEADER);
        String roleType = headers.getFirst(ROLE_TYPE_HEADER);
        if(!ObjectUtils.isEmpty(userId)&&!ObjectUtils.isEmpty(roleType)){
            ThreadLocalUtils.set(ThreadAttributes.getInstance(Long.parseLong(userId), Integer.parseInt(roleType)));
        }
        // 子协议设置
        List<String> protocols = request.getHeaders().get(SEC_WEBSOCKET_PROTOCOL);
        if (protocols != null && !protocols.isEmpty()) {
            String protocol = protocols.get(0); // 取第一个子协议
            response.getHeaders().put(SEC_WEBSOCKET_PROTOCOL, List.of(protocol)); // 返回给客户端
        }
        return true;
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response,
                               WebSocketHandler wsHandler, Exception exception) {
    }
}
